open Sast
open Ast
open Type
open Operator

module SymTbl = Map.Make(struct
  type t = string
  let compare x y = Pervasives.compare x y
end)


(* (from type, to type)  *)
let rec checkType types = match types with 
    | (_, Int) -> 
        (match types with
         | (Int, _)|(Float, _)|(RInt, _)|(RFloat,_) -> Int
         | (_, _)   -> failwith "Type Convert Mismatch (Int)"
         )
    | (_, Float) ->
        (match types with
         | (Int, _)|(Float, _)|(RInt, _)|(RFloat,_) -> Float
         | (_, _)   -> failwith "Type Convert Mismatch (Float)"
         )
    | (_, RInt) -> 
        (match types with
         | (RInt, _) -> RInt
         | (_, _)   -> failwith "Type Convert Mismatch (RInt)"
         )
    | (_ , RFloat) -> 
        (match types with
         | (RFloat, _) -> RFloat
         | (_, _)   -> failwith "Type Convert Mismatch (RFloat)"
         )
    | (_ , Tuple) -> Tuple
    | (_, Str) ->
        (match types with
         | (Str, _)|(Int, _)|(Float, _)|(RInt, _)|(RFloat,_)|(Vector 0,_) -> Str
         | (_, _)   -> failwith "Type Convert Mismatch (Str)"
         )
    | (Vector 0, Vector 0) -> Vector 0
    | (_, Func) -> Func
    | (_, _)   -> failwith "Type Convert Mismatch"
;;

let rec canConversion types = match types with 
    | (Int, _)|(RInt, _) -> 
        (match types with
         | (_, Int)|(_, RInt) -> Int
         | (_, Float)|(_, RFloat) -> Float
         | (_, Str) -> Str
         | (_, _)   -> failwith "If/Else Mismatch (Int/RInt)"
         )
    | (Float, _)|(RFloat, _) -> 
        (match types with
         | (_, Int)|(_, RInt) -> Float
         | (_, Float)|(_, RFloat) -> Float
         | (_, Str) -> Str
         | (_, _)   -> failwith "If/Else Mismatch (Float)"
         )
    | (Str, _) ->
        (match types with
         | (_, Str)|(_, Int)|(_, RInt)|(_, Float)|(_, RFloat)|(_, Vector 0) -> Str
         | (_, _)   -> failwith "If/Else Mismatch (Str)"
         )
    | (Vector 0, Vector 0) -> Vector 0
    | (Vector 0, Str) -> Str
    | (Tuple, Tuple) -> Tuple
    | (_, _) -> failwith "Cannot convert"
;;

let convert_arg_type to_args from_args = 
    let (v, typ) = from_args in (from_args, (checkType (typ, to_args)))
;;

(* find function and return a return type *)
let rec findfun id = function
    | [] ->  
        if Builtin.exists id 
            then Builtin.get_ret_type id
            else failwith ("Undeclared Function:" ^ id)
    | loc::scp ->
        let (vars,funs) = loc in
        if SymTbl.mem id funs 
            then 
                let (typ, arg_typ) = SymTbl.find id funs
                in typ
            else findfun id scp
;;

(* find a function from name and return argument types *)
let rec findfunargs id = function
    | [] ->  
        if Builtin.exists id 
            then Builtin.get_arg_types id
            else failwith ("Undeclared Function:" ^ id)
    | loc::scp ->
        let (vars,funs) = loc in
        if SymTbl.mem id funs 
            then 
                let (typ, arg_typ) = SymTbl.find id funs
                in arg_typ
            else findfunargs id scp
;;


(* find a function and return variable types *)
let rec findvar id = function
    | [] -> failwith ("Undeclared Variable:" ^ id)
    | loc::scp ->
        let (vars,funs) = loc in
        if SymTbl.mem id vars 
            then SymTbl.find id vars
            else if SymTbl.mem id funs
                then let (typ, args) = SymTbl.find id funs
                    in Func (* typ *)               (* TODO:a type of function variables may be function... *)
                else findvar id scp
;;

let chk ast = 
    let rec expr env = function
              Ast.LInt e -> Sast.LInt e, Int 
            | Ast.LFloat e -> Sast.LFloat e, Float
            | Ast.Binop(e1, op, e2) ->
                let e1 = expr env e1
                and e2 = expr env e2 in
                let _, t1 = e1
                and _, t2 = e2
                in (match op with
                |  Add -> (match (t1,t2) with 
                    | (Int, _) | (RInt, _) -> (match (t1,t2) with
                        | (_, Int) | (_, RInt)      -> Sast.Binop(e1, op, e2), Int
                        | (_, Float) | (_, RFloat)  -> Sast.Binop(e1, op, e2), Float
                        | (_, Str)  -> Sast.Binop(e1, op, e2), Str
                        | (_, _)                    -> failwith "Unsupported Add")
                    | (Float, _) | (RFloat, _) -> (match (t1,t2) with
                        | (_, Int) | (_, RInt)|(_, Float) | (_, RFloat) -> Sast.Binop(e1, op, e2), Float
                        | (_, Str)  -> Sast.Binop(e1, op, e2), Str
                        | (_, _)                    -> failwith "Unsupported Add")
                    | (Vector a, _) -> (match (t1,t2) with
                        | (_, Vector b) -> Sast.Binop(e1, op, e2), Vector a
                        | (_, Str)  -> Sast.Binop(e1, op, e2), Str
                        | (_, _)                    -> failwith "Unsupported Add")
                    | (Str, _) -> (match (t1,t2) with
                        | (_, Int) | (_, RInt)|(_, Float) | (_, RFloat)|(_, Str)|(_, Vector 0) -> Sast.Binop(e1, op, e2), Str
                        | (_, _)                    -> failwith "Unsupported Add")
                    | (_, _) -> failwith "Unsupported Add/Sub")

                |  Sub -> (match (t1,t2) with 
                    | (Int, _) | (RInt, _) -> (match (t1,t2) with
                        | (_, Int) | (_, RInt)      -> Sast.Binop(e1, op, e2), Int
                        | (_, Float) | (_, RFloat)  -> Sast.Binop(e1, op, e2), Float
                        | (_, _)                    -> failwith "Unsupported Sub")
                    | (Float, _) | (RFloat, _) -> (match (t1,t2) with
                        | (_, Int) | (_, RInt)|(_, Float) | (_, RFloat) -> Sast.Binop(e1, op, e2), Float
                        | (_, _)                    -> failwith "Unsupported Sub")
                    | (Vector a, _) -> (match (t1,t2) with
                        | (_, Vector b) -> Sast.Binop(e1, op, e2), Vector a
                        | (_, _)                    -> failwith "Unsupported Sub")
                    | (_, _) -> failwith "Unsupported Add/Sub")

                | Equal|Neq|Greater|Geq|Leq|Less|Land|Lor -> 
                (* use Int instead of boolean *)
                (match (t1,t2) with 
                     | (Int, _) | (RInt, _)| (Float, _) | (RFloat, _) -> (match (t1,t2) with
                         | (_, Int) | (_, RInt)| (_, Float) | (_, RFloat) -> Sast.Binop(e1, op, e2), Int
                         | (_, _) -> failwith "Unsupported Compare")
                     | (_, _) -> failwith "Unsupported Compare")

                | Mult -> (match (t1, t2) with
                     | (Int, _) | (RInt, _) -> (match (t1,t2) with
                         | (_, Int) | (_, RInt)      -> Sast.Binop(e1, op, e2), Int
                         | (_, Float) | (_, RFloat)  -> Sast.Binop(e1, op, e2), Float
                         | (_, Vector a)               -> Sast.Binop(e1, op, e2), Vector a
                         | (_, _)                    -> failwith "Unsupported Mult")
                     | (Float, _) | (RFloat, _) -> (match (t1,t2) with
                         | (_, Int) | (_, RInt)|(_, Float) | (_, RFloat) -> Sast.Binop(e1, op, e2), Float
                         | (_, _)                    -> failwith "Unsupported Mult")
                         | (_, Vector a)               -> Sast.Binop(e1, op, e2), Vector a
                     | (_, _) -> failwith "Unsupported Mult")

                | Div -> (match (t1, t2) with
                     | (Int, _) | (RInt, _) -> (match (t1,t2) with
                         | (_, Int) | (_, RInt)      -> Sast.Binop(e1, op, e2), Int
                         | (_, Float) | (_, RFloat)  -> Sast.Binop(e1, op, e2), Float
                         | (_, _)                    -> failwith "Unsupported Div")
                     | (Float, _) | (RFloat, _) -> (match (t1,t2) with
                         | (_, Int) | (_, RInt)      -> Sast.Binop(e1, op, e2), Float
                         | (_, Float) | (_, RFloat)  -> Sast.Binop(e1, op, e2), Float
                         | (_, _)                    -> failwith "Unsupported Div")
                     | (_, _) -> failwith "Unsupported Div")

                | Rmndr -> (match (t1, t2) with
                     | (Int, _) | (RInt, _) -> (match (t1,t2) with
                         | (_, Int) | (_, RInt)      -> Sast.Binop(e1, op, e2), Int
                         | (_, _)                    -> failwith "Unsupported Reminder")
                     | (_, _) -> failwith "Unsupported Reminder")

                | Innrp -> (match (t1,t2) with 
                     | (Vector a, Vector b)  -> Sast.Binop(e1, op, e2), Float
                     | (_, _) -> failwith "Unsupported Innrp")
                )

            | Ast.Uop(op, e) -> 
                            let e = expr env e in
                            let _, t = e in
                            (match op with
                            | Minus -> (match t with
                            | Int|Float|RInt|RFloat|(Vector _) -> Sast.Uop(op, e), t
                            | _ -> failwith "Unsupported Minus")
                            
                            (* Right now, Not could be operated on any type
                             * Everything could be regarded as boolean 1 expect 0*) 
                            | Not -> Sast.Uop(op, e), t)
           | Ast.LString str -> Sast.LString str, Str
           | Ast.LList e -> 
                           Sast.LList (List.fold_left (fun ls ex -> (expr env ex)::ls)
                           [] e), List
           | Ast.LVctr e -> Sast.LVctr (List.fold_left (fun ls ex -> (expr env ex)::ls)
                           [] e), Vector (List.length e)
           | Ast.LTple e -> Sast.LTple (List.fold_left (fun ls ex -> (expr env ex)::ls)
                           [] e), Tuple
           | Ast.Elmt (s, e) ->
                           let e = expr env e in Sast.Elmt (s, e), Int
           | Ast.Id e -> Sast.Id e, (findvar e env)

           | Ast.Call (func_name, args) ->
               let sast_args = List.fold_left(fun ls ex -> ls@[(expr env ex)]) [] args in          (* Types and value in call *)
               let fun_args =  (findfunargs func_name env) in                                      (* Types in decrlation *)
               let cnvt_new_args = (List.map2 convert_arg_type fun_args sast_args) in 
               Sast.Call(func_name, cnvt_new_args), findfun func_name env
               
           | Ast.If(e, e1, e2) ->
                    let e = expr env e in
                    let _, t = e in
                    (match t with
                        | Int -> 
                            let (eif, tpif) = (expr env e1) in
                            let (eelse, tpelse) = (expr env e2) in
                            let typ = canConversion(tpif, tpelse) in
                                Sast.If(e, (eif, tpif), (eelse, tpelse)), typ
                        | _ -> failwith "Predicate of if must be integer"
                    )

           | Ast.Scope (inits, body) -> 
                    let env = (SymTbl.empty, SymTbl.empty)::env 
                    in
                    let (s1' , env)= 
                        List.fold_left 
                            (fun (st, env') s -> 
                                let (stmt, env) = (stmt env' s)
                                in
                                (st@[stmt], env)
                            ) 
                            ([], env) inits
                    in
                    let (body, typ) = (expr env body) 
                    in
                    Sast.Scope (s1', (body, typ)), typ

    and
    stmt env = function

           | Ast.VDclr(t, str, e) -> (match env with
                [] -> failwith "empty scope in FDclr"
                | (lvars, lfuns)::globals -> 
                    let lvars = SymTbl.add str t lvars in 
                    let env = (lvars, lfuns)::globals in (
                Sast.VDclr(t, str, expr env e), env))
           | Ast.While (s, e) -> failwith "While is not supported!"
           | Ast.FDclr (t1, str, args, e) -> (match env with
                [] -> failwith "empty scope in FDclr"
                | (lvars, lfuns)::globals -> 
                    let args_types = List.fold_left (fun acc (typ, id) -> acc@[typ]) [] args in
                    let lfuns = SymTbl.add str (t1, args_types) lfuns in 
                    let lvars_new = List.fold_left (fun scope' (typ, var) -> (SymTbl.add var typ scope')) SymTbl.empty args in
                    let env = [(lvars_new, SymTbl.empty); (lvars, lfuns)]@globals in (
                Sast.FDclr(t1, str, args, expr env e), env))
    in

    let emptyScope = [(SymTbl.empty, SymTbl.empty)] in

    fst (List.fold_left (fun (prog, env) st -> let stmt' = (stmt env st) in prog@[(fst stmt')], snd(stmt')) ([], emptyScope) ast)

